//
// Created by kuang on 2024/11/25.
//

#include "pcap_utils.h"


// 内存操作函数实现
uint32_t write_pcap_global_header_to_buffer(uint8_t *buffer) {
    struct pcap_global_header header = {
            .magic_number = 0xa1b2c3d4,
            .version_major = 2,
            .version_minor = 4,
            .thiszone = 0,
            .sigfigs = 0,
            .snaplen = 65535,
            .network = 1
    };

    memcpy(buffer, &header, sizeof(header));
    return sizeof(header);
}

uint32_t write_udp_packet_to_buffer(uint8_t *buffer,
                                    const uint8_t *src_mac, const uint8_t *dst_mac,
                                    uint32_t src_ip, uint32_t dst_ip,
                                    uint16_t src_port, uint16_t dst_port,
                                    const uint8_t *payload, uint16_t payload_length) {
    struct pcap_packet_header pkt_header = {0};
    struct ethernet_header eth_header = {0};
    struct ip_header ip_hdr = {0};
    struct udp_header udp_hdr = {0};
    uint32_t offset = 0;

    // 设置以太网头部
    memcpy(eth_header.dst_mac, dst_mac, 6);
    memcpy(eth_header.src_mac, src_mac, 6);
    eth_header.ethertype = htons(0x0800);

    // 设置IP头部
    ip_hdr.version_ihl = 0x45;
    ip_hdr.tos = 0;
    ip_hdr.total_length = htons(20 + 8 + payload_length);
    ip_hdr.identification = htons(0);
    ip_hdr.flags_fragment_offset = htons(0x4000);
    ip_hdr.ttl = 64;
    ip_hdr.protocol = 17;
    ip_hdr.src_ip = src_ip;
    ip_hdr.dst_ip = dst_ip;
    ip_hdr.header_checksum = 0;
    ip_hdr.header_checksum = calculate_checksum((uint16_t *)&ip_hdr, sizeof(ip_hdr));

    // 设置UDP头部
    udp_hdr.src_port = htons(src_port);
    udp_hdr.dst_port = htons(dst_port);
    udp_hdr.length = htons(8 + payload_length);
    udp_hdr.checksum = 0;

    // 设置PCAP包头
    time_t current_time;
    time(&current_time);
    pkt_header.ts_sec = (uint32_t)current_time;
    pkt_header.ts_usec = 0;
    pkt_header.incl_len = sizeof(eth_header) + sizeof(ip_hdr) + sizeof(udp_hdr) + payload_length;
    pkt_header.orig_len = pkt_header.incl_len;

    // 写入所有数据
    memcpy(buffer + offset, &pkt_header, sizeof(pkt_header));
    offset += sizeof(pkt_header);
    memcpy(buffer + offset, &eth_header, sizeof(eth_header));
    offset += sizeof(eth_header);
    memcpy(buffer + offset, &ip_hdr, sizeof(ip_hdr));
    offset += sizeof(ip_hdr);
    memcpy(buffer + offset, &udp_hdr, sizeof(udp_hdr));
    offset += sizeof(udp_hdr);
    memcpy(buffer + offset, payload, payload_length);
    offset += payload_length;

    return offset;
}

// 文件操作函数实现
void write_pcap_global_header_to_file(FILE *fp) {
    uint8_t buffer[sizeof(struct pcap_global_header)];
    uint32_t size = write_pcap_global_header_to_buffer(buffer);
    fwrite(buffer, 1, size, fp);
}

void write_udp_packet_to_file(FILE *fp,
                              const uint8_t *src_mac, const uint8_t *dst_mac,
                              uint32_t src_ip, uint32_t dst_ip,
                              uint16_t src_port, uint16_t dst_port,
                              const uint8_t *payload, uint16_t payload_length) {
    uint8_t buffer[65535];
    uint32_t size = write_udp_packet_to_buffer(buffer, src_mac, dst_mac,
                                               src_ip, dst_ip, src_port, dst_port,
                                               payload, payload_length);
    fwrite(buffer, 1, size, fp);
}

// 解析函数实现
void parse_pcap_buffer(const uint8_t *buffer, uint32_t buffer_length) {
    uint32_t offset = 0;

    // 解析PCAP全局头
    struct pcap_global_header *global_header = (struct pcap_global_header *)buffer;
    printf("\n=== PCAP Global Header ===\n");
    printf("Magic Number: 0x%08x\n", global_header->magic_number);
    printf("Version: %d.%d\n", global_header->version_major, global_header->version_minor);
    printf("Snaplen: %u\n", global_header->snaplen);
    printf("Network: %u\n", global_header->network);

    offset += sizeof(struct pcap_global_header);

    // 解析数据包
    while (offset < buffer_length) {
        // 解析PCAP包头
        struct pcap_packet_header *pkt_header = (struct pcap_packet_header *)(buffer + offset);
        printf("\n=== Packet Header ===\n");
        printf("Timestamp: %u.%u\n", pkt_header->ts_sec, pkt_header->ts_usec);
        printf("Captured Length: %u\n", pkt_header->incl_len);
        printf("Actual Length: %u\n", pkt_header->orig_len);

        offset += sizeof(struct pcap_packet_header);

        // 解析以太网头
        struct ethernet_header *eth_header = (struct ethernet_header *)(buffer + offset);
        char mac_str[18];
        printf("\n=== Ethernet Header ===\n");
        printf("Destination MAC: %s\n", mac_addr_to_str(eth_header->dst_mac, mac_str, sizeof(mac_str)));
        printf("Source MAC: %s\n", mac_addr_to_str(eth_header->src_mac, mac_str, sizeof(mac_str)));
        printf("EtherType: 0x%04x\n", ntohs(eth_header->ethertype));

        offset += sizeof(struct ethernet_header);

        // 解析IP头
        struct ip_header *ip_hdr = (struct ip_header *)(buffer + offset);
        char ip_str[16];
        printf("\n=== IP Header ===\n");
        printf("Version: %d\n", (ip_hdr->version_ihl >> 4) & 0xF);
        printf("Header Length: %d bytes\n", (ip_hdr->version_ihl & 0xF) * 4);
        printf("Total Length: %d\n", ntohs(ip_hdr->total_length));
        printf("TTL: %d\n", ip_hdr->ttl);
        printf("Protocol: %d\n", ip_hdr->protocol);
        printf("Source IP: %s\n", ip_addr_to_str(ip_hdr->src_ip, ip_str, sizeof(ip_str)));
        printf("Destination IP: %s\n", ip_addr_to_str(ip_hdr->dst_ip, ip_str, sizeof(ip_str)));

        offset += sizeof(struct ip_header);

        // 解析UDP头
        struct udp_header *udp_hdr = (struct udp_header *)(buffer + offset);
        printf("\n=== UDP Header ===\n");
        printf("Source Port: %d\n", ntohs(udp_hdr->src_port));
        printf("Destination Port: %d\n", ntohs(udp_hdr->dst_port));
        printf("Length: %d\n", ntohs(udp_hdr->length));

        offset += sizeof(struct udp_header);

        // 显示payload
        uint16_t payload_length = ntohs(udp_hdr->length) - sizeof(struct udp_header);
        printf("\n=== Payload ===\n");
        printf("Length: %d bytes\n", payload_length);
        print_hex_dump(buffer + offset, payload_length);

        offset += payload_length;
    }
}

// 辅助函数实现
uint16_t calculate_checksum(uint16_t *addr, int len) {
    int count = len;
    uint32_t sum = 0;
    uint16_t answer = 0;

    while (count > 1) {
        sum += *addr++;
        count -= 2;
    }

    if (count > 0) {
        sum += *(uint8_t *)addr;
    }

    while (sum >> 16) {
        sum = (sum & 0xffff) + (sum >> 16);
    }

    answer = ~sum;
    return answer;
}

void print_hex_dump(const uint8_t *data, int len) {
    for (int i = 0; i < len; i++) {
        if (i % 16 == 0) printf("\n%04x: ", i);
        printf("%02x ", data[i]);
        if ((i + 1) % 16 == 0) {
            printf("  ");
            for (int j = i - 15; j <= i; j++) {
                if (data[j] >= 32 && data[j] <= 126)
                    printf("%c", data[j]);
                else
                    printf(".");
            }
        }
    }
    printf("\n");
}

char* mac_addr_to_str(const uint8_t *mac_addr, char *str, size_t size) {
    snprintf(str, size, "%02x:%02x:%02x:%02x:%02x:%02x",
             mac_addr[0], mac_addr[1], mac_addr[2],
             mac_addr[3], mac_addr[4], mac_addr[5]);
    return str;
}

char* ip_addr_to_str(uint32_t ip_addr, char *str, size_t size) {
    struct in_addr addr;
    addr.s_addr = ip_addr;
    strncpy(str, inet_ntoa(addr), size);
    return str;
}